package com.coolw.data.migrate.common.util;

import com.coolw.data.migrate.common.SpringContextHolder;
import com.coolw.data.migrate.common.constant.SqlTemplateConst;
import com.coolw.data.migrate.common.exception.JdbcException;
import lombok.extern.slf4j.Slf4j;
import org.dromara.hutool.core.collection.CollUtil;
import org.dromara.hutool.core.lang.Assert;
import org.dromara.hutool.core.text.StrUtil;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Arrays;
import java.util.List;

/**
 * JDBC工具类
 *
 * @author coolw
 * @date 2024/2/1 16:09
 */
@Slf4j
@Component
public class JdbcUtil implements DisposableBean {

    private static DataSource dataSource;

    private static final ThreadLocal<Connection> CONN_THREAD_LOCAL = new ThreadLocal<>();

    /**
     * 获取表数据大小
     *
     * @param tableName 表名
     * @return 数据大小
     */
    public static int getTableDataSizeByName(String tableName) {
        String sql = StrUtil.format(SqlTemplateConst.COUNT_SQL_TEMPLATE, tableName);
        log.info("getTableDataSizeByName,tableName:{},sql:{}", tableName, sql);
        Assert.notBlank(tableName, "tableName 不能为空");

        Connection connection = null;
        Statement statement = null;
        ResultSet resultSet = null;
        try {
            connection = getConnection();

            statement = connection.createStatement();

            resultSet = statement.executeQuery(sql);
            if (resultSet.next()) {
                return resultSet.getInt(1);
            }
            return 0;
        } catch (Exception e) {
            log.error("getTableDataSizeByName error,tableName:{}", tableName, e);
            throw new JdbcException("getTableDataSizeByName error");
        } finally {
            //close(connection, statement, resultSet);
        }
    }

    /**
     * 获取表数据大小
     *
     * @param sql 完整的sql语句
     * @return 数据大小
     */
    public static int getTableDataSizeBySql(String sql) {
        log.info("getTableDataSizeBySql,sql:{}", sql);
        Assert.notBlank(sql, "sql 不能为空");

        Connection connection = null;
        Statement statement = null;
        ResultSet resultSet = null;
        try {
            connection = getConnection();

            statement = connection.createStatement();

            resultSet = statement.executeQuery(sql);
            if (resultSet.next()) {
                return resultSet.getInt(1);
            }
            return 0;
        } catch (Exception e) {
            log.error("getTableDataSizeBySql error,sql:{}", sql, e);
            throw new JdbcException("getTableDataSizeByName error");
        } finally {
            //close(connection, statement, resultSet);
        }
    }

    /**
     * SQL执行,手动控制事务
     *
     * @param insertSql 新增SQL
     * @param deleteSql 删除SQL
     * @return 执行结果
     */
    public static boolean sqlExecuteWithManualTra(String insertSql, String deleteSql) {
        return batchSqlExecuteWithManualTra(Arrays.asList(insertSql, deleteSql));
    }

    /**
     * 批量SQL执行,手动控制事务
     *
     * @param sqlList sql语句列表
     * @return 执行结果
     */
    public static boolean batchSqlExecuteWithManualTra(List<String> sqlList) {
        log.info("batchSqlExecuteWithManualTra,sqlList:{}", sqlList);
        if (CollUtil.isEmpty(sqlList)) {
            return false;
        }

        Connection connection = null;
        Statement statement = null;
        try {
            connection = getConnection();
            connection.setAutoCommit(false);

            statement = connection.createStatement();
            for (String sql : sqlList) {
                statement.addBatch(sql);
            }
            statement.executeBatch();

            connection.commit();
            return true;
        } catch (Exception e) {
            log.error("batchSqlExecuteWithManualTra error,sqlList:{},", sqlList, e);
            rollback(connection);
            return false;
        } finally {
            //close(connection, statement, null);
        }
    }

    public static Connection getConnection() {
        Connection conn = CONN_THREAD_LOCAL.get();
        if (conn != null) {
            return conn;
        }

        dataSource = SpringContextHolder.getBean(DataSource.class);
        try {
            conn = dataSource.getConnection();
        } catch (SQLException e) {
            throw new JdbcException("Data source connection failed", e);
        }

        CONN_THREAD_LOCAL.set(conn);
        return conn;
    }

    private static void rollback(Connection connection) {
        if (connection == null) {
            log.warn("Connection is null");
            return;
        }
        try {
            connection.rollback();
        } catch (SQLException e) {
            log.error("Rollback error", e);
        }
    }

    public static void clear() {
        Connection conn = CONN_THREAD_LOCAL.get();
        if (conn != null) {
            close(conn, null, null);
        }
        CONN_THREAD_LOCAL.remove();
    }

    private static void close(Connection conn, Statement stat, ResultSet rs) {
        if (rs != null) {
            try {
                rs.close();
            } catch (SQLException e) {
                log.error("ResultSet close error", e);
            }
        }

        if (stat != null) {
            try {
                stat.close();
            } catch (SQLException e) {
                log.error("Statement close error", e);
            }
        }

        if (conn != null) {
            try {
                conn.close();
            } catch (SQLException e) {
                log.error("Connection close error", e);
            }
        }
    }

    @Override
    public void destroy() throws Exception {
        if (log.isDebugEnabled()) {
            log.debug("Data source destroy......");
        }
        dataSource = null;

        clear();
    }
}
